package com.authentication.code2Token;

import com.authentication.code.AuthenticationCodeRepository;
import com.authentication.code.ClientAndUserAuthenticationPair;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.authentication.BadCredentialsException;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.*;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2RefreshTokenGenerator;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;

import java.time.Duration;
import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.temporal.TemporalAmount;
import java.util.*;

import static com.authentication.filter.CodeStateAuthenticationFilter.USER_GRANTED_AUTHORITIES;

public class Code2TokenAuthenticationManager implements AuthenticationManager {

    private OAuth2TokenGenerator<Jwt> oAuth2TokenGenerator;

    private OAuth2RefreshTokenGenerator oAuth2RefreshTokenGenerator;
    private final RefreshTokenRepository refreshTokenRepository;

    private final AuthenticationCodeRepository authenticationCodeRepository;

    public Code2TokenAuthenticationManager(AuthenticationCodeRepository authenticationCodeRepository, RefreshTokenRepository refreshTokenRepository) {
        this.authenticationCodeRepository = authenticationCodeRepository;
        this.refreshTokenRepository = refreshTokenRepository;

    }

    @Autowired
    public void setOAuth2RefreshTokenGenerator(OAuth2RefreshTokenGenerator oAuth2RefreshTokenGenerator) {
        this.oAuth2RefreshTokenGenerator = oAuth2RefreshTokenGenerator;
    }


    @Autowired
    public void setOAuth2TokenGenerator(OAuth2TokenGenerator<Jwt> oAuth2TokenGenerator) {
        this.oAuth2TokenGenerator = oAuth2TokenGenerator;
    }


    @Override
    public OAuth2AccessTokenAuthenticationToken authenticate(Authentication authentication) throws AuthenticationException {

        Authentication userPrincipal = null;

        Set<String> userGrantedAuthorities = null;

        if (authentication instanceof OAuth2AuthorizationCodeAuthenticationToken authorizationCodeToken) {
//          使用 Authorization code 换取 token

            ClientAndUserAuthenticationPair clientAndUserAuthenticationPair = authenticationCodeRepository.get(authorizationCodeToken.getCode());

            userPrincipal = clientAndUserAuthenticationPair.getUserAuthentication();

            userGrantedAuthorities = new HashSet<>(
                    Arrays.asList(
                            clientAndUserAuthenticationPair.getGrantedScopes()
                    )
            );

        } else if (authentication instanceof OAuth2RefreshTokenAuthenticationToken refreshToken) {
//            使用刷新令牌 换取 token
            OAuth2TokenContext oAuth2TokenContext = refreshTokenRepository.get(refreshToken.getRefreshToken());

            if (oAuth2TokenContext == null) {
                throw new BadCredentialsException("无效的 refresh token");
            }

            userPrincipal = oAuth2TokenContext.getAuthorizationGrant();

            userGrantedAuthorities = oAuth2TokenContext.getAuthorizedScopes();

        } else {
            throw new UnsupportedOperationException("暂不支持换取 token");
        }

        OAuth2ClientAuthenticationToken clientAuthenticationToken = (OAuth2ClientAuthenticationToken) SecurityContextHolder.getContext().getAuthentication();

        RegisteredClient registeredClient = clientAuthenticationToken.getRegisteredClient();

        if (registeredClient == null) {
            throw new IllegalStateException();
        }


        OAuth2AccessToken accessToken = createAccessToken(registeredClient, userPrincipal, userGrantedAuthorities);

        OAuth2RefreshToken oAuth2RefreshToken = createRefreshToken(registeredClient, userPrincipal, userGrantedAuthorities);

        HashMap<String, Object> additionalParams = new HashMap<>();

        additionalParams.put("refreshTokenExpiresAt", Date.from(oAuth2RefreshToken.getExpiresAt()));


        return new OAuth2AccessTokenAuthenticationToken(
                registeredClient,
                clientAuthenticationToken,
                accessToken,
                oAuth2RefreshToken,
                additionalParams
        );

    }

    /**
     * 创建刷新令牌
     *
     * @param registeredClient
     * @param userPrincipal
     * @param userGrantedAuthorities
     * @return
     */
    private OAuth2RefreshToken createRefreshToken(RegisteredClient registeredClient, Authentication userPrincipal, Set<String> userGrantedAuthorities) {

        DefaultOAuth2TokenContext tokenContext = DefaultOAuth2TokenContext.builder()
                .tokenType(OAuth2TokenType.REFRESH_TOKEN)
                .authorizationGrant(userPrincipal)
                .registeredClient(registeredClient)
                .authorizedScopes(userGrantedAuthorities)
                .build();


        OAuth2RefreshToken refreshToken = oAuth2RefreshTokenGenerator.generate(tokenContext);

        if (refreshToken == null || refreshToken.getTokenValue() == null) {
            throw new IllegalStateException("创建 refreshToken 失败");
        }

//      把 refreshToken 临时持久化
        refreshTokenRepository.save(refreshToken.getTokenValue(), tokenContext);
        return refreshToken;


    }

    private OAuth2AccessToken createAccessToken(RegisteredClient registeredClient, Authentication userPrincipal, Set<String> userGrantedAuthorities) {

        Instant now = nowInstantPlusMinutes(Duration.ZERO);

        Duration accessTokenTimeToLive = registeredClient.getTokenSettings().getAccessTokenTimeToLive();

        Instant expiredTime = nowInstantPlusMinutes(accessTokenTimeToLive);

        //        读取 用户授予 client 的权限列表

        DefaultOAuth2TokenContext tokenContext = DefaultOAuth2TokenContext.builder()

                .authorizationGrantType(AuthorizationGrantType.AUTHORIZATION_CODE)
                .registeredClient(registeredClient)
                .principal(userPrincipal)

                // 令牌持有用户授予的权限,而非 client 注册是可以请求的权限
                .authorizedScopes(userGrantedAuthorities)
                .tokenType(OAuth2TokenType.ACCESS_TOKEN)
                .build();

        Jwt jwt = oAuth2TokenGenerator.generate(tokenContext);


        return new OAuth2AccessToken(
                OAuth2AccessToken.TokenType.BEARER,
                jwt.getTokenValue(),
                now,
                expiredTime,
                userGrantedAuthorities
        );

    }


    /**
     * @param minutes 要加上多少分钟
     * @return 当前时间 加上 对应分钟之后的 instant
     */
    private static Instant nowInstantPlusMinutes(TemporalAmount minutes) {
        return LocalDateTime.now().plus(minutes).atZone(ZoneId.systemDefault()).toInstant();
    }
}
